# data_classification.py

import os
import sys
import random
import shutil

def check_folders(destination_folder):
    required_folders = [
        os.path.join(destination_folder, 'images', 'test'),
        os.path.join(destination_folder, 'images', 'val'),
        os.path.join(destination_folder, 'labels', 'val')
    ]

    for folder in required_folders:
        if not os.path.exists(folder):
            print(f"文件夹不存在: {folder}")
            return False

    print("所有必需的文件夹都存在。")
    return True

def retain_images(destination_folder):
    images_test_folder = os.path.join(destination_folder, 'images', 'test')
    test_images = os.listdir(images_test_folder)

    # 随机删除图片，只保留 30%
    num_to_keep = int(len(test_images) * 0.3)
    images_to_keep = random.sample(test_images, num_to_keep)

    # 删除不在保留列表中的图片
    for image in test_images:
        if image not in images_to_keep:
            os.remove(os.path.join(images_test_folder, image))
            print(f"删除图片: {image}")

def create_temp_folder(destination_folder):
    temp_val_folder = os.path.join(destination_folder, 'newsval')
    os.makedirs(temp_val_folder, exist_ok=True)
    return temp_val_folder

def process_images_and_labels(destination_folder, temp_val_folder):
    images_val_folder = os.path.join(destination_folder, 'images', 'val')
    labels_val_folder = os.path.join(destination_folder, 'labels', 'val')

    # 随机选择 40% 的图像
    val_images = os.listdir(images_val_folder)
    num_val_images = int(len(val_images) * 0.4)
    selected_val_images = random.sample(val_images, num_val_images)

    # 移动选中的图像和所有标签文件
    for image in selected_val_images:
        shutil.move(os.path.join(images_val_folder, image), temp_val_folder)

    for label in os.listdir(labels_val_folder):
        shutil.move(os.path.join(labels_val_folder, label), temp_val_folder)

def clear_original_folders(destination_folder):
    folders_to_clear = [
        os.path.join(destination_folder, 'images', 'val'),
        os.path.join(destination_folder, 'labels', 'val')
    ]

    for folder in folders_to_clear:
        for file in os.listdir(folder):
            os.remove(os.path.join(folder, file))
        print(f"已清空文件夹: {folder}")

def compare_and_remove_unmatched_files(temp_val_folder):
    val_images = set(os.listdir(temp_val_folder))
    val_labels = set(os.listdir(temp_val_folder))

    for label in val_labels:
        if label != 'classes.txt' and label.replace('.txt', '.jpg') not in val_images:
            os.remove(os.path.join(temp_val_folder, label))
            print(f"删除不匹配的标签文件: {label}")

def move_files_back(destination_folder, temp_val_folder):
    # 移动 newsval 中的文件回原始文件夹
    for file in os.listdir(temp_val_folder):
        if file.endswith('.jpg'):
            target_folder = os.path.join(destination_folder, 'images', 'val')
        elif file.endswith('.txt'):
            target_folder = os.path.join(destination_folder, 'labels', 'val')
        else:
            continue  # 跳过其他文件类型

        # 确保目标文件夹存在
        os.makedirs(target_folder, exist_ok=True)

        # 打印目标路径
        print(f"移动文件: {file} 到 {target_folder}")

        # 移动文件
        shutil.move(os.path.join(temp_val_folder, file), target_folder)

def delete_temp_folder(temp_val_folder):
    shutil.rmtree(temp_val_folder)
    print("已删除临时文件夹 newsval。")

def main(destination_folder):
    print(f"目标文件夹路径: {destination_folder}")

    try:
        if check_folders(destination_folder):
            retain_images(destination_folder)
            temp_val_folder = create_temp_folder(destination_folder)
            process_images_and_labels(destination_folder, temp_val_folder)
            clear_original_folders(destination_folder)# 清空原始文件夹
            compare_and_remove_unmatched_files(temp_val_folder)
            move_files_back(destination_folder, temp_val_folder)
            delete_temp_folder(temp_val_folder)
        else:
            print("请确保所有必需的文件夹都存在。")
    except Exception as e:
        print(f"发生错误: {e}")

if __name__ == "__main__":
    if len(sys.argv) > 1:
        destination_folder = sys.argv[1]
    else:
        destination_folder = ""

    main(destination_folder)